#!/usr/bin/env python
# -*- coding: utf-8 -*-
import json
import time
import traceback
import threading
import websocket
import logging


class WebsocketClient(object):
    def __init__(self, callback_handler, callback_in_thread=True, url="ws://127.0.0.1:9000/ws_client"):
        self.callback_handler = callback_handler
        self.callback_in_thread = callback_in_thread
        self.url = url
        self.is_open = False
        self.exit = False

    def run(self):
        while not self.exit:
            self.ws = websocket.WebSocketApp(url=self.url,
                                             on_message=self.on_message,
                                             on_error=self.on_error,
                                             on_close=self.on_close,
                                             on_open=self.on_open)
            try:
                self.ws.run_forever()
            except:
                pass
            time.sleep(1)

    def stop(self):
        logging.info("websocket_client stop")
        self.exit = True
        self.ws.close()

    def start(self):
        logging.info("websocket_client start")
        if not self.exit:
            self.is_open = False
            threading.Thread(target=self.run).start()

    def restart(self):
        logging.info("websocket_client restart")
        self.is_open = False
        self.ws.close()

    def on_message(self, ws, msg):
        logging.info("websocket_client on_message: " + msg)
        j_msg = json.loads(msg)
        if "key" in j_msg:
            callback = getattr(self.callback_handler, j_msg["key"], None)
            if callback:
                try:
                    if self.callback_in_thread:
                        threading.Thread(target=callback, args=[j_msg["data"]], daemon=True).start()
                    else:
                        callback(j_msg["data"])
                except:
                    traceback.print_exc()

    def on_error(self, ws, error):
        logging.info("websocket_client on_error{}\n{}", type(error), error)
        self.is_open = False

    def on_close(self, ws):
        logging.info("websocket_client on_close")
        self.is_open = False

    def on_open(self, ws):
        logging.info("websocket_client on_open")
        self.is_open = True

    def write_message(self, key, data=None):
        if not self.is_open:
            return
        if len(key) == 0:
            return
        if data is None:
            data = {}
        try:
            msg = {"key": key, "data": data}
            msg_str = json.dumps(msg)
            logging.info("websocket_client write_message: " + msg_str)
            self.ws.send(msg_str)
        except:
            traceback.print_exc()
